'''
Created on Jan 27, 2012

@author: alebalbin
'''
import os
import sys
import numpy as np
from collections import defaultdict

class MicroArrSet:
    
    def __init__(self,genNames, FeatureNames, file_names,exp,pvals):
        self.genes=genNames
        self.features=FeatureNames
        self.exp=-1*exp # This is because AJim prepares microarrays, so Agilent software generates -log10
        self.pvals=pvals
        self.exp2=self.log2transform()
        self.filenames=file_names
    
    def log2transform(self):
        '''
        np.log2(1000) = np.log10(1000)/np.log10(2)
        '''
        mat = self.exp/np.log10(2)
        return mat
        
    
    def find_repeated_features(self,features):
        '''
        '''
        ids=defaultdict(list)
        for i, g in features.iteritems():
            ids[g].append(i)
        return ids
    
    def summaryze_repeat_features(self,log2=True):
        '''
        summaryze repeated features using the features with highest sum across samples
        '''
        if log2:
            mat=self.exp2
        else:
            mat=self.exp
        gids=self.find_repeated_features(self.genes)
        nmat = np.empty( (len(gids), mat.shape[1]) )
        ngids=defaultdict(list)
        i=0
        for f,rows in gids.iteritems():
            rows=np.array(rows)
            '''
            #to get the row with max sum
            s = np.sum(mat[rows,],axis=1)
            s = rows[s==np.max(s)]
            if len(s)>1:
                s=s[0]
            nmat[i,:]=mat[s,:]
            '''
            nmat[i,:]=np.mean(mat[rows,],axis=0) #to get the mean summary            
            ngids[f].append(i)
            i+=1
        return nmat, ngids
    
    
def count_lines(file):
    ifile = open(file)
    Data=False
    
    for l in ifile:
        fields=l.strip('\n').split('\t')
        if fields[0]=='FEATURES':
            hd=fields
            Data=True
            ind=0
            ft=['chr_coord','ProbeName','GeneName', 'SystematicName', 'Description', 'LogRatio', 'LogRatioError', 'PValueLogRatio']
            
            d=defaultdict()
            print file, hd
            for f in ft:
                d[f]=hd.index(f)
        if Data:
            ind+=1
    ifile.close()
    return ind,d

def read_files(rootfolder,file_list):
    fdict=defaultdict()
    for n,f in file_list.iteritems():
        nrows,d=count_lines(os.path.join(rootfolder,f))
        # Check that all files have the same size
    cols=1
    nfiles=len(file_list)
    exp=np.empty((nrows,cols*nfiles))
    pval=np.empty((nrows,cols*nfiles))
    ft=['chr_coord','ProbeName','GeneName', 'SystematicName', 'Description']
    gnames=defaultdict(list)
    features=defaultdict()
    fnames=defaultdict()
    j=0
    for fname,file in file_list.iteritems():
        ifile = open(os.path.join(rootfolder,file))
        Data=False
        names=defaultdict()     
        gnames=defaultdict(list)
        for l in ifile:
            fields=l.strip('\n').split('\t')
            if not Data and fields[0]=='FEATURES':
                ind=0
                Data=True
                continue
            if Data:
                if fields[d['Description']]=="Unknown" or fields[d['Description']]=="":
                    continue
                
                exp[ind,j*cols:j*cols+cols]=np.array(map(float,[fields[d['LogRatio']]] ))                    
                pval[ind,j*cols:j*cols+cols]=np.array(map(float,[fields[d['PValueLogRatio']]] ))
                
                names[ind]=[fields[d['chr_coord']],fields[d['ProbeName']],fields[d['GeneName']], 
                            fields[d['SystematicName']], fields[d['Description']]]
                gnames[ind]=fields[d['GeneName']] #.append(fields[d['GeneName']])
                
                ind+=1
                        
        #fname=file.split('_')[-1]
        features[j]=(fname,names)
        fnames[j]=fname
        j+=1
        ifile.close()
        
    return MicroArrSet(gnames,features,fnames,exp,pval)
    

def write_microarrays(genids,exp,ofile,filenames):
    '''
    '''
    of=open(ofile,'w')
    hd=["NAME"]+filenames.values()
    print hd, filenames
    of.write(",".join(hd).replace(',','\t')+'\n')
    for f,i in genids.iteritems():
        of.write(f+"\t"+",".join(map(str,list(exp[i,:][0])) ).replace(',','\t')+'\n')
    of.close()

def read_sample_table(file):
    ifile=open(file)
    file_list=defaultdict()
    for l in ifile:
        f = l.strip('\n').split('\t')
        print f
        file_list[f[0]]=f[1]
    ifile.close()
    return file_list
        
'''
files=['/Users/alebalbin/Desktop/Dropbox/KRAS/LCK/GE_data_H441/07262011_251485068639_S01_GE2-v5_95_Feb07_1_1_H441-LCK9.txt',
       '/Users/alebalbin/Desktop/Dropbox/KRAS/LCK/GE_data_H441/07262011_251485068639_S01_GE2-v5_95_Feb07_1_2_H441-LCK11.txt',
       '/Users/alebalbin/Desktop/Dropbox/KRAS/LCK/GE_data_H441/07262011_251485068639_S01_GE2-v5_95_Feb07_1_3_H441-LCK12.txt'
       ]

ofile='/Users/alebalbin/Desktop/Dropbox/KRAS/LCK/GE_data_H441/H441-LCK_aggregate.txt'
'''
'''
files=['/Users/alebalbin/Desktop/Dropbox/KRAS/LCK/GE_data_H358/data_8-11-11/08112011_251485068643_S01_GE2-v5_95_Feb07_1_3_H358-LCK9.txt',
       '/Users/alebalbin/Desktop/Dropbox/KRAS/LCK/GE_data_H358/data_8-11-11/08112011_251485068643_S01_GE2-v5_95_Feb07_1_4_H358-LCK11.txt',
       '/Users/alebalbin/Desktop/Dropbox/KRAS/LCK/GE_data_H358/data_8-11-11/08112011_251485068644_S01_GE2-v5_95_Feb07_1_1_H358-LCK12.txt']

ofile='/Users/alebalbin/Desktop/Dropbox/KRAS/LCK/GE_data_H358/H358-LCK_aggregate.txt'
'''

#sample_table='/Users/alebalbin/Desktop/Dropbox/KRAS/Control_CellLines/lck_controls_sample_table.txt'
#ofile='/Users/alebalbin/Desktop/Dropbox/KRAS/Control_CellLines/lck_celllines_control_aggregate.txt'

#sample_table='/Users/alebalbin/Desktop/Dropbox/KRAS/LCK/GE_data_H358/H358_MET_sample_table.txt'
#ofile='/Users/alebalbin/Desktop/Dropbox/KRAS/LCK/enrichment_analysis/H358_H441_analysis/H358_aggregate.txt'
#rootfolder='/Users/alebalbin/Desktop/Dropbox/KRAS/LCK/GE_data_H358/data_8-11-11/'
'''
sample_table='/Users/alebalbin/Desktop/Dropbox/KRAS/LCK/GE_data_H441/H441_MET_sample_table.txt'
ofile='/Users/alebalbin/Desktop/Dropbox/KRAS/LCK/enrichment_analysis/H358_H441_analysis/H441_MET_aggregate.txt'
rootfolder='/Users/alebalbin/Desktop/Dropbox/KRAS/LCK/GE_data_H441'
'''

# Bushra analysis
sample_table='/Users/alebalbin/Documents/projects/forCollaboration/forBushra/Bushra_Microarray_data/spink_sample_table.txt'
ofile='/Users/alebalbin/Documents/projects/forCollaboration/forBushra/Bushra_Microarray_data/SPINK_KD_aggregate.txt'
rootfolder='/Users/alebalbin/Documents/projects/forCollaboration/forBushra/Bushra_Microarray_data/'


files=read_sample_table(sample_table)



# This script throw out probes with unknown annotation
# Summarizes probes by gene, and compute the average expression for all the probes that map to a gene
# It also assumes that you need to take the -log10. that it is the way Jim prepare the microarrays
microArrays = read_files(rootfolder,files)
exp,genids = microArrays.summaryze_repeat_features()
write_microarrays(genids,exp,ofile,microArrays.filenames)
        
